# $Id: snmp.py 7557b5d4d82b 2012/12/02 09:07:16 pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#
# essential imports
__version__ = "$Revision: 7557b5d4d82b $"
__author__ = "$Author: pgurumur $"

from pysnmp.carrier.asynsock.dispatch import AsynsockDispatcher
from pysnmp.carrier.asynsock.dgram import udp
from pyasn1.codec.ber import encoder, decoder
from pysnmp.proto import api
from pysnmp.smi import builder, view, error
from time import time

from djanet.lib.exception import SNMPError

GET = 1
SET = 2
WALK = 3

class snmp(object):
   def __init__(self, **Keywords):
      self._err = SNMPError()

      self._host = None
      self._comm = None
      self._port = 0
      self._snmp = None
      self._type = None
      self._reqmesg = None
      self._conn = False

      if "server" in Keywords:
         self._host = Keywords["server"].strip()

      if "community" in Keywords:
         self._comm = Keywords["community"].strip()

      if "port" in Keywords:
         try:
            self._port = int(Keywords["port"])
         except ValueError, message:
            raise ValueError, message

      if self._port == 0:
         self._port = 161

      if "version" in Keywords:
         if Keywords["version"] == "1":
            self._snmp = api.protoModules[api.protoVersion1]
         elif Keywords["version"] == "2c":
            self._snmp = api.protoModules[api.protoVersion2c]
         elif Keywords["version"] == "2":
            self._snmp = api.protoModules[api.protoVersion2c]

      if not self._snmp:
         self._snmp = api.protoModules[api.protoVersion2c]

      self._builder = builder.MibBuilder().loadModules()
      self._view = view.MibViewController(self._builder)
      self._dispatch = AsynsockDispatcher()
      self._started = time()
      self._retval = {}

   _getid = property(fget = lambda self: self._snmp.apiPDU.getRequestID)
   value = property(fget = lambda self: self._retval)
   close = property(fget = lambda self: self._dispatch.closeDispatcher())
   integer = property(fget = lambda self: self._snmp.Integer)
   integer32 = property(fget = lambda self: self._snmp.Integer32)
   ipaddress = property(fget = lambda self: self._snmp.IpAddress)
   null = property(fget = lambda self: self._snmp.Null)
   string = property(fget = lambda self: self._snmp.OctetString)
   timer = property(fget = lambda self: self._snmp.TimeTicks)
   unsigned = property(fget = lambda self: self._snmp.Unsigned32)
   opaque = property(fget = lambda self: self._snmp.Opaque)

   def _error(self, Message):
      if Message:
         raise SNMPError, "%s: %s" %(str(self), str(Message))

   def __str__(self):
      return str("%s: %d" %(self._host, self._port)).encode('utf-8')

   def __repr__(self):
      return repr(str(self))

   def __unicode__(self):
      return unicode(str(self))

   def __len__(self):
      return len(self._retval.values())

   def _setpdu(self, snmptype):
      if snmptype == GET:
         self._reqpdu = self._snmp.GetRequestPDU()
      elif snmptype == SET:
         self._reqpdu = self._snmp.SetRequestPDU()
      elif snmptype == WALK:
         self._reqpdu = self._snmp.GetNextRequestPDU()

      self._snmp.apiPDU.setDefaults(self._reqpdu)
      self._type = snmptype

   def _timer(self, now, started = time()):
      if (now - started) > 3:
         self._error('request timed out!')

   def _transport(self):
      self._retval = {}
      if not self._conn:
         self._reqmesg = self._snmp.Message()
         self._snmp.apiMessage.setDefaults(self._reqmesg)
         self._snmp.apiMessage.setCommunity(self._reqmesg, self._comm)
         self._conn = True

      self._snmp.apiMessage.setPDU(self._reqmesg, self._reqpdu)
      self._dispatch.registerTransport(udp.domainName,
            udp.UdpSocketTransport().openClientMode())
      self._dispatch.registerRecvCbFun(self._recvCallBack)
      self._dispatch.sendMessage(encoder.encode(self._reqmesg),
            udp.domainName, (self._host, self._port))
      self._dispatch.jobStarted(True)
      self._dispatch.runDispatcher()
      self._dispatch.closeDispatcher()

   def _recvCallBack(self, dispatcher, domain, address, message, reqpdu = None):
      if not reqpdu:
         reqpdu = self._reqpdu

      while message:
         response, message = decoder.decode(message,
               asn1Spec = self._snmp.Message())
         respdu = self._snmp.apiMessage.getPDU(response)
         if self._getid(reqpdu) == self._getid(respdu):
            status = self._snmp.apiPDU.getErrorStatus(respdu)
            if (self._type == GET) or (self._type == SET):
               if status:
                  self._error(status.prettyPrint())

               for oid, val in self._snmp.apiPDU.getVarBinds(respdu):
                  self._retval[oid.prettyPrint()] = val.prettyPrint()
            
               self._dispatch.jobFinished(True)
            elif self._type == WALK:
               if status and (status != 2):
                  self._error(status.prettyPrint())

               varbind = self._snmp.apiPDU.getVarBindTable(reqpdu, respdu)
               for row in varbind:
                  for name, val in row:
                     if val is None:
                        continue

                     self._retval[name.prettyPrint()] = val.prettyPrint()

               for oid, val in varbind[-1]:
                  if val is not None:
                     break

               else:
                  self._dispatch.jobFinished(True)

               self._snmp.apiPDU.setVarBinds(
                     reqpdu, map(lambda (ix, jy), val = self._snmp.Null(): (ix,
                        val), varbind[-1]))

               self._snmp.apiPDU.setRequestID(reqpdu,
                     self._snmp.getNextRequestID())
               self._dispatch.sendMessage(
                     encoder.encode(self._reqmesg), domain, address)

               if (time() - self._started) > 3:
                  self._error('request timed out')

               self._started = time()

      return message

   def Get(self, Oid):
      self._setpdu(GET)
      self._snmp.apiPDU.setVarBinds(self._reqpdu, ((Oid, self._snmp.Null()),))
      self._transport()
      return self._retval.values()

   def Set(self, Oid, Value):
      self._setpdu(SET)
      self._snmp.apiPDU.setVarBinds(self._reqpdu, ((Oid, Value), ))
      self._transport()
      return self._retval.values()

   def Walk(self, Oid):
      self._setpdu(WALK)
      headVars = [ self._snmp.ObjectIdentifier((Oid)) ]
      self._snmp.apiPDU.setVarBinds(self._reqpdu, map(
         lambda ix, snmp = self._snmp: (ix, self._snmp.Null()), headVars))
      self._transport()
      return self._retval
